import numpy as np
import SimpleITK as sitk

def rand(*args):
    """ 
    获得随机数 
    if len(args) == 1 return [0, args[0]]
    if len(args) == 2 return [args[0], args[1]]
    """
    if len(args) == 1:
        return 0 if args[0] < 0 else int(np.random.rand() * args[0])
    elif len(args) == 2:
        if args[1] < args[0]:
            raise Exception("Invalid args!")
        else:
            return int(np.random.rand() * (args[1] - args[0] + 1) + args[0])
    else:
        raise Exception("Invalid args!")

def normalize(image, low=700.0, high=3000.0):
    """ 归一化灰度值 """
    normalized = (image - low) / (high - low)
    return np.clip(normalized, 0.0, 1.0)

def distance_transform(label):
    """ 距离变换 """
    # sitk.SignedDanielssonDistanceMapImageFilter is more precise but more slowly
    label = label.astype(np.uint8)
    distance_filter = sitk.SignedMaurerDistanceMapImageFilter()
    distance_filter.SetSquaredDistance(False)
    distance_filter.SetUseImageSpacing(False)
    ret = distance_filter.Execute(sitk.GetImageFromArray(label))
    return sitk.GetArrayFromImage(ret)

def gauss_smoothing(array):
    """高斯平滑"""
    smoothing_filter = sitk.SmoothingRecursiveGaussianImageFilter()
    smoothing_filter.SetSigma(0.4)
    smoothing_filter.SetNumberOfThreads(8)
    return sitk.GetArrayFromImage(smoothing_filter.Execute(sitk.GetImageFromArray(array)))

def gauss_noising(array):
    """高斯噪声"""
    noise_filter = sitk.AdditiveGaussianNoiseImageFilter()
    noise_filter.SetMean(rand(20, 25))
    noise_filter.SetNumberOfThreads(8)
    noise_filter.SetStandardDeviation(rand(20, 25))
    return sitk.GetArrayFromImage(noise_filter.Execute(sitk.GetImageFromArray(array)))

def flip_image(images, axis):
    """翻转图像"""
    return [np.flip(image, axis=axis) for image in images]

def label_smoothing(label):
    """平滑标签"""
    smoothing_filter = sitk.SmoothingRecursiveGaussianImageFilter()
    smoothing_filter.SetSigma(0.25)
    smoothing_filter.SetNormalizeAcrossScale(True)
    smoothing_filter.SetNumberOfThreads(8)
    sitk_label = sitk.GetImageFromArray(label)
    smoothed = sitk.GetArrayFromImage(smoothing_filter.Execute(sitk_label))
    return np.clip(smoothed, 0.0, 1.0)

def intensityTranslation(image):
    """灰度平移"""
    return image + rand(-20, 20)

def rotate_image(images):
    """随机旋转图片"""
    def _rotate(image, rotation, interpolator):
        image = sitk.GetImageFromArray(image)
        resample_filter = sitk.ResampleImageFilter()
        resample_filter.SetInterpolator(interpolator)
        resample_filter.SetTransform(rotation)
        resample_filter.SetReferenceImage(image)
        resampled = resample_filter.Execute(image)
        return sitk.GetArrayFromImage(resampled)

    size = inputs[0].shape
    center = [size[2 - i] / 2.0 for i in range(3)]
    angle = rand(-5, 5) / 57.3
    rotation = sitk.Euler3DTransform()
    rotation.SetRotation(*angle)
    rotation.SetCenter(center)

    outputs = []
    for image in images:
        output = _rotate(image, rotation, sitk.sitkLinear)
        outputs.append(np.around(output))
    return outputs

'''
def elastix_deformation(images, deformation=4, num_grid=2):
    """ 随机拉伸变换 """
    itk_image = sitk.GetImageFromArray(images[0])
    bspline_transformation = sitk.BSplineTransformInitializer(
        itk_image, [num_grid] * itk_image.GetDimension())
    params = bspline_transformation.GetParameters()
    params = np.asarray(params, dtype=float)
    params += np.random.randn(params.shape[0]) * deformation
    bspline_transformation.SetParameters(tuple(params))

    resampler = sitk.ResampleImageFilter()
    resampler.SetReferenceImage(itk_image)
    resampler.SetInterpolator(sitk.sitkLinear)
    resampler.SetDefaultPixelValue(0)
    resampler.SetTransform(bspline_transformation)

    outputs = []
    for image in images:
        output = resampler.Execute(
            sitk.GetImageFromArray(image.astype(np.float32)))
        output = sitk.GetArrayFromImage(output)
        output = np.round(output).astype(np.int16)
        outputs.append(output)
    return outputs
'''

def write_image(file_name, image):
    writer = sitk.ImageFileWriter()
    writer.SetFileName(file_name)
    writer.Execute(image)

def read_image(file_name):
    reader = sitk.ImageFileReader()
    reader.SetFileName(file_name)
    return reader.Execute()

def slide(array, starts, sizes):
    slices = [slice(start, start + size) for start, size in zip(starts, sizes)]
    return array[slices]
